﻿#pragma once

enum Color {
	RED,
	BLACK
};

template <class K, class V>
struct RBTreeNode {
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	pair<K, V> _kv;
	Color _col;

	RBTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
		, _col(RED) {
	}
};

template <class K, class V>
struct RBTree {
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv) {
		if (_root == nullptr) { //如果插入节点是根
			_root = new Node(kv);
			_root->_col = BLACK;
			return true;
		}
		//不是根，找插入位置。记录插入节点的父节点方便链接
		Node* parent = nullptr, * cur = _root;
		while (cur) {
			if (kv.first > cur->_kv.first) {
				parent = cur;
				cur = cur->_right;
			}
			else if (kv.first < cur->_kv.first) {
				parent = cur;
				cur = cur->_left;
			}
			else
				return false;
		}
		//新增黑色节点会影响所有路径
		//新增红色节点只会影响父节点
		cur = new Node(kv);
		cur->_col = RED;
		//插入节点链接在其父节点的左或右
		if (kv.first > parent->_kv.first) {
			parent->_right = cur;
			cur->_parent = parent;
		}
		else {
			parent->_left = cur;
			cur->_parent = parent;
		}
		//检查是否需要变色
		//cur的父节点为空(即根节点)，或其父节点不是红色就不用更新
		while (parent && parent->_col == RED) {
			Node* grandfather = parent->_parent;
			if (parent == grandfather->_left) { //通过g找到u。p为g的左孩子
				Node* uncle = grandfather->_right;
				//u存在且为红(cur为红，p为红，g为黑)
				if (uncle && uncle->_col == RED) { //u存在且为红
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					//继续往上更新，去找上一个红节点
					cur = grandfather;//grandfather有可能为根，所以父节点可能不存在，所以需要在循环条件中判断
					parent = cur->_parent;
				}
				else { //u不存在或u为黑(cur为红，p为红，g为黑)
					//p为g的左孩子，cur为p的左孩子，则进行右单旋转。(这里p已经是g的左)
					if (cur == parent->_left) {
						RotateR(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else { //p为g的左孩子，cur为p的右孩子。(这里p已经是g的左)
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
			else { //parent == grandfather->_right;
				Node* uncle = grandfather->_left;
				if (uncle && uncle->_col == RED) {
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else {
					if (cur == parent->_right) {
						RotateL(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else { //cur == parent->_left
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
		}
		//根节点必须为黑，所以不用在上面的过程中考虑grandfather为根时怎么变色的问题
		_root->_col = BLACK;
		return true;
	}

	void RotateL(Node* parent) {
		//parent是平衡因子为2的节点
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		Node* parentParent = parent->_parent;//parent不一定是整棵树的根
		//旋转完后，subR是子树的根，还需要处理subR的parent

		//旋转
		parent->_right = subRL;
		subR->_left = parent;
		//上方还未处理完，根节点没有确定，各节点的parent还未更新
		parent->_parent = subR;
		if (subRL) subRL->_parent = parent;

		//处理根节点
		if (_root == parent) { //parent是整棵树的根
			_root = subR;
			subR->_parent = nullptr;
		}
		else { //parent不一定是整棵树的根，旋转完后，还需要处理subR的parent
			if (parent == parentParent->_left) parentParent->_left = subR;
			else parentParent->_right = subR;
			subR->_parent = parentParent;
		}
	}

	void RotateR(Node* parent) {
		Node* subL = parent->_left, * subLR = subL->_right, * parentParent = parent->_parent;

		//旋转
		parent->_left = subLR;
		subL->_right = parent;
		//处理各节点的父节点
		parent->_parent = subL;
		if (subLR) subLR->_parent = parent;

		//处理根节点
		if (_root == parent) {
			_root = subL;
			subL->_parent = nullptr;
		}
		else {
			if (parent == parentParent->_left) parentParent->_left = subL;
			else parentParent->_right = subL;
			subL->_parent = parentParent;
		}
	}

	Node* Find(const K& key) {
		Node* cur = _root;
		while (cur) {
			if (key > _root->_kv.first) cur = cur->_right;
			else if (key < _root->_kv.first) cur = cur->_left;
			else return cur;
		}
		return nullptr;
	}

	void InOrder() {
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance() {
		if (_root == nullptr) return true;
		if (_root->_col == RED) return false;
		int blacknum = 0;//根节点到当前节点的黑色节点数量
		//blacknum只能传值，不能传引用。引用求得是整棵树所有黑色节点

		//将最左路径节点作为参考值传入Check和其他路径比较
		int refVal = 0;
		Node* cur = _root;
		while (cur) {
			if (cur->_col == BLACK) ++refVal;
			cur = cur->_left;
		}
		return Check(_root, blacknum, refVal);
	}

	int Height() { return _Height(_root); }
	size_t Size() { return _Size(_root); }
private:
	Node* _root = nullptr;

	void _InOrder(Node* root) {
		if (root == nullptr)
			return;
		_InOrder(root->_left);
		cout << root->_kv.first << " ";
		_InOrder(root->_right);
	}

	bool Check(Node* root, int blacknum, const int refVal) {
		if (root == nullptr) {
			//cout << blacknum << endl;
			//走到空意味着求出了一条路径的黑色节点数量
			if (refVal != blacknum) {
				cout << "存在黑色节点数量不相等的路径" << endl;
				return false;
			}
			return true;
		}
		if (root->_col == RED && root->_parent->_col == RED) {
			cout << "有连续的红色节点" << endl;
			return false;
		}
		if (root->_col == BLACK) ++blacknum;
		return Check(root->_left, blacknum, refVal)
			&& Check(root->_right, blacknum, refVal);
	}

	int _Height(Node* root) {
		if (root == nullptr) return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	size_t _Size(Node* root) {
		if (root == nullptr) return 0;
		return _Size(root->_left) + _Size(root->_right) + 1;
	}
};